'''
chi_detector.py
---
locate chinese word in image.
'''
import sys
sys.path = ['/home/hzh/.local/lib/python3.5/site-packages/'] + sys.path

import os
from os.path import join
import cv2
import yaml
import numpy as np
from matplotlib import pyplot as plt
import shutil


def x_histogram(img):
    return np.sum(img, 1)

def y_histogram(img):
    return np.sum(img, (0, 2))

def plot_hist_of_img(img, hist):
    print("hist shape=", hist.shape)
    plt.subplot(211)
    plt.plot(range(len(hist)), hist, 'ro-')
    plt.subplot(212)
    # plt.imshow(img)
    # plt.show()


def search_bars(img, hist):
    Thres = 1000
    bars = [[0, 0]]
    NewBar_GAP = 2
    for x in range(len(hist)-1):
        if hist[x] < Thres:
            if x - bars[-1][1] > NewBar_GAP:
                bars.append([x, x])
            else:
                bars[-1][1] = x
    return bars

def search_bbox_seq(bars, h):
    prev = None
    WIDTH_LOW = 50
    WIDTH_HIGH = 65
    boxes = []
    for b in bars[h:]:
        if prev == None:
            prev = b
        d = b[0] - prev[1]
        # print(prev, "--", b, " | ", d)
        if WIDTH_HIGH < d:
            # print("new")
            prev = b
        elif WIDTH_LOW < d:
            # print("append")
            boxes.append((prev[1], b[0]))
            prev = b
    return boxes

def detect_bbox(bars):
    max_score = -1
    best_split = None
    for h in range(min(3, len(bars)-1)):
        bbox_seq = search_bbox_seq(bars, h)
        if len(bbox_seq) > max_score:
            max_score = len(bbox_seq)
            best_split = bbox_seq
    # widths = []
    # for a, b in zip(bars[:-1], bars[1:]):
        # w = b[0] - a[1] 
        # widths.append(w)
    # print("widths:", widths)
    return best_split

def plot_bar(img, bars, bbox):
    for b in bars:
        img[:, b] = [np.random.randint(0, 255) for _ in range(3)]
    rows = img.shape[0]
    for x in bbox:
        img[:, (x[0]+1, x[1]-1)] = (0, 255, 0)
        img[10, np.arange(*x)] = (0, 255, 0)
        img[rows-10, np.arange(*x)] = (0, 255, 0)
    plt.imshow(img)
    plt.show()


def init_db(db_dir):
    if os.path.exists(db_dir):
        cmd = input('%s already exist. Delete? y/[n]' % db_dir)
        if cmd in 'yY':
            shutil.rmtree(db_dir)
        else:
            exit(-1)
    os.makedirs(db_dir)

def remove_margin(img):
    rows = img.shape[0]
    top, down = 0, rows-1
    for t in range(0, rows):
        s = np.sum(img[t, :])
        if(s > 1000):
            top = t
            break
    for d in range(rows-1, top, -1):
        s = np.sum(img[d, :])
        if(s > 1000):
            down = d
            break
    return img[top:down+1, :]


class FontDetector:

    def __init__(self, in_dir, out_dir):
        self._input_dir = in_dir
        self._db_dir = out_dir
        init_db(self._db_dir)
        if not os.path.exists(in_dir):
            raise Exception('no such input db dir:%s' % in_dir)
        with open(join(in_dir, 'index.yaml'), 'r') as f:
            self._frame_dict = yaml.load(f.read())
        self._index = {}

    def save(self):
        print("dumping index file...")
        with open(join(self._db_dir, 'index.yaml'), 'w') as f:
            f.write(yaml.dump(self._index))
        print("dump completed.")

    def save_chi_imgs(self, i, img, bbox):
        print(i)
        self._index[i] = {'time': self._frame_dict[i],
                          'chi_seq': []}
        for k, b in enumerate(bbox):
            fn = '%d_%d.png' % (i, k)
            fp = join(self._db_dir, fn)
            im = img[10:-10, np.arange(b[0]-1, b[1]+1)]
            im = remove_margin(im)
            cv2.imwrite(fp, im)
            self._index[i]['chi_seq'].append(fn)
        print(i, ':', self._index[i]['chi_seq'])

    def run(self):
        # cnt = 0
        for i in self._frame_dict.keys():
            img_fp = join(self._input_dir, '%d.png' % i)
            img = cv2.imread(img_fp)
            yh = y_histogram(img)
            # print('y hist shape=', yh.shape)
            # plot_hist_of_img(img, yh)
            bars = search_bars(img, yh)
            # print('bars=', bars)
            bbox = detect_bbox(bars)
            # print('bbox=', bbox)
            # plot_bar(img, bars, bbox)
            self.save_chi_imgs(i, img, bbox)
            # cnt += 1
            # if cnt > 50:
                # break
        self.save()



if __name__ == '__main__':
    d = FontDetector("./bottom_imgs", "./chi_seq_imgs")
    d.run()
